#! /usr/bin/env python
# coding=utf-8

# 根据 文件名 修改 直系同源基因的id号
# 如果 可选 nucl 则对应生成一个文件 为对应的核酸序列
# 注意 nucl中的":"会被替换为"_"

import os
import shutil
from Bio import SeqIO
import Bio
import sys
import argparse


parser = argparse.ArgumentParser(
    description='''将筛选完的直系同源群 改名 并建立一个核酸文件与之对应（可选）
    用法:
    change_ortho_name_and_match_nucl.py -i Orthogroup_Sequences_filtered -n nucl_total.fna -o Orthogroup_Sequences_namechanged -O Orthogroup_Sequences_namechanged_nucl
    由大天才于2021年7月19日创建于浙江农业大学''')

parser.add_argument('-i',
                help='必须给定，Orthogroup_Sequences_filtered的路径')

parser.add_argument('-n',
                help='可选, 对应的核酸序列文件')


parser.add_argument('-o',
                help='必须给定，输出AA文件的路径')

parser.add_argument('-O',
                help='当-n选择是必须给定，输出nucl文件的路径')


args = parser.parse_args()

match_nucl  = True

if not args.n:

    match_nucl  = False

else:
    infile_nucl = args.n




if not args.i or not args.o:

    parser.print_help()

    sys.exit()

if match_nucl == True:

    if not args.O:

        parser.print_help()

        sys.exit()

    else:
        outfile_nucl = args.O 





infile = args.i 
outfile = args.o





try:
    shutil.rmtree(outfile)
except:
    pass

try:
    os.mkdir(outfile)
except:
    pass





nucl_dic = {}


if match_nucl:

    try:
        shutil.rmtree(outfile_nucl)
    except:
        pass

    try:
        os.mkdir(outfile_nucl)
    except:
        pass

    for i in SeqIO.parse(infile_nucl,'fasta'):

        # 这里 用_ 替换 : 是因为 orthofinder 也这么替换 

        nucl_dic[str(i.name).replace(':','_')] = str(i.seq)



for ogs_file in os.listdir(infile):

    out_handel = open(outfile+'/'+ogs_file,'w')

    used_lista_len = {}

    if match_nucl:
        out_nucl_handel = open(outfile_nucl+'/'+ogs_file,'w')

        nucl_len = {}
    

    for j in SeqIO.parse(infile+'/'+ogs_file,'fasta'):

        spe = j.name.split('@')[0]

        name = spe+'@'+ogs_file.split('.')[0]

        if match_nucl:
        
            if j.name not in nucl_dic:
                
                raise BaseException('错误， 缺少匹配核酸序列 %s' % j.name)

        if name not in used_lista_len:

            used_lista_len[name] = str(j.seq)

            if match_nucl:

                nucl_len[name] = nucl_dic[j.name]
        else:
            if len(str(j.seq)) > len(used_lista_len[name]):

                used_lista_len[name] = str(j.seq)

                if match_nucl:

                    nucl_len[name] = nucl_dic[j.name]




    for i in used_lista_len:

        out_handel.write('>'+i+'\n'+used_lista_len[i]+'\n')

        if match_nucl:

            out_nucl_handel.write('>'+i+'\n'+nucl_len[i]+'\n')

            

        

    out_handel.close()
    if match_nucl:
        out_nucl_handel.close()
        

